from datetime import datetime, time as datetime_time, timedelta
import time
import sys
import os

# The connectivity from one city to another
class Connectivity:
	def __init__(self, dept_city, arrv_city, min_durat, num_trains, valid=True):
		self.dept_city = dept_city
		self.arrv_city = arrv_city
		self.min_durat = min_durat		# Type: `datetime_time`
		self.num_trains = num_trains	# Type: integer
		self.valid = valid				# Type: boolean

# Reduced train information class
class TrainInfo_r:
	def __init__(self, train_no, dept_city, arrv_city, duration, op_days):
		self.train_no = train_no		# Type: string
		self.dept_city = dept_city		# Type: string (of city name)
		self.arrv_city = arrv_city
		self.duration = duration		# Type: `datetime_time`
		self.op_days = op_days			# Type: integer

# Read the detailed CSV file and return a list of `TrainInfo_s` objects
def read_reduced_csv(file_name):
	train_list = []
	# Use UTF-8 encoding to read the file
	with open(file_name, "r", encoding="utf-8") as file:
		lines = file.readlines()
		# Skip the first line
		lines = lines[1:]
		for line in lines:
			train_info = line.split(",")
			train_no = train_info[0]
			dept_city = train_info[1]
			arrv_city = train_info[2]
			duration = datetime.strptime(train_info[3], "%H:%M").time()
			op_days = int(train_info[4])
			train_list.append(TrainInfo_r(train_no, dept_city, arrv_city, duration, op_days))
	return train_list

# Summarize the connectivity between cities and return a list of `Connectivity` objects
def summarize(data_path):
	# Initialize the connectivity list
	connectivity_list = []
	# Read every CSV file in the directory
	for file_name in os.listdir(data_path):
		if file_name.endswith("-r.csv"):
			file_name = os.path.join(data_path, file_name)
			train_list = read_reduced_csv(file_name)
			for train in train_list:
				# Search the connectivity in the list
				flag_found = False
				for conn in connectivity_list:
					if conn.dept_city == train.dept_city and conn.arrv_city == train.arrv_city:
						flag_found = True
						if conn.min_durat > train.duration:
							conn.min_durat = train.duration
						conn.num_trains += train.op_days
						break
				if not flag_found:
					# If not found, append a new connectivity
					connectivity_list.append(Connectivity(train.dept_city, train.arrv_city, train.duration, train.op_days))
	return connectivity_list

# Merge the connectivity list to undirected graph
# A connectivity from B to A will be merged to the connectivity from A to B,
# the `min_durat` and `num_trains` will be the geometric mean of the two connectivities
def merge(connectivity_list):
	# Initialize the merged connectivity list
	merged_list = []
	# Merge the connectivity list
	for conn in connectivity_list:
		flag_found = False
		for merged in merged_list:
			if conn.dept_city == merged.arrv_city and conn.arrv_city == merged.dept_city:
				flag_found = True
				# Calculate the geometric mean of `min_durat`
				min_durat_1 = conn.min_durat.hour * 60 + conn.min_durat.minute
				min_durat_2 = merged.min_durat.hour * 60 + merged.min_durat.minute
				min_durat_g = int((min_durat_1 * min_durat_2) ** 0.5 + 0.5)
				merged.min_durat = datetime_time(min_durat_g // 60, min_durat_g % 60)
				# Calculate the geometric mean of `num_trains`
				merged.num_trains = int((conn.num_trains * merged.num_trains) ** 0.5 + 0.5)
				# Set the merged connectivity to valid
				# If one of the connectivity is not found, the merged connectivity will be invalid
				merged.valid = True
				break
		if not flag_found:
			# If not found, append a new connectivity
			merged_list.append(Connectivity(conn.dept_city, conn.arrv_city, conn.min_durat, conn.num_trains, False))
	return merged_list

def main():
	# Read command line arguments
	if len(sys.argv) < 2:
		print("Usage: python DataSummary.py <data_path>")
		exit(1)
	connectivity_list = summarize(sys.argv[1])
	# Sort the connectivity list by the departure city
	connectivity_list.sort(key=lambda x: x.dept_city)
	# Merge the connectivity list to undirected graph
	connectivity_list_undirected = merge(connectivity_list)
	# Write the connectivity list to a CSV file
	output_file = os.path.join(sys.argv[1], "connectivity.csv")
	with open(output_file, "w", encoding="utf-8") as f:
		f.write("Departure City,Arrival City,Minimum Duration,Number of Trains\n")
		for conn in connectivity_list:
			f.write(f"{conn.dept_city},{conn.arrv_city},{conn.min_durat.strftime('%H:%M')},{conn.num_trains}\n")
	print(f"Data has been summarized to {output_file}.")
	# Write the undirected connectivity list to a CSV file
	output_file_undirected = os.path.join(sys.argv[1], "connectivity-undirected.csv")
	with open(output_file_undirected, "w", encoding="utf-8") as f:
		f.write("City 1,City 2,Minimum Duration,Number of Trains,Valid\n")
		for conn in connectivity_list_undirected:
			f.write(f"{conn.dept_city},{conn.arrv_city},{conn.min_durat.strftime('%H:%M')},{conn.num_trains},{"1" if conn.valid else "0"}\n")

if __name__ == "__main__":
	main()
